#include "HitsEvent.h"
#include "KalmanFilterProcess.h"
#include "CalculationKit.h"
#include <TRandom3.h>

KalmanFilterProcess::KalmanFilterProcess()
{

}

HitsEvent KalmanFilterProcess::filter_2D(HitsEvent input_hitsEvent)
{	
	vector<float> fx;
    vector<float> fy;
    vector<float> fz;
    vector<float> fe;
	HitsEvent KF_result(fx, fy, fz, fe);
	return KF_result;
}

HitsEvent KalmanFilterProcess::filter_3D(HitsEvent input_hitsEvent)
{
    typedef ROOT::Math::SMatrix<double,6,6>  SMatrix66;
	typedef ROOT::Math::SMatrix<double,3,6>  SMatrix36;
	typedef ROOT::Math::SMatrix<double,6,3>  SMatrix63;
	typedef ROOT::Math::SMatrix<double,3,3>  SMatrix33;
	typedef ROOT::Math::SVector<double,6>  SVector6;
	typedef ROOT::Math::SVector<double,3>  SVector3;
	typedef ROOT::Math::SVector<double,3>  SVector3;

	double MeasureError = 0.3;
	double MeasureErrorx = 0.3;
	double MeasureErrory = 0.3;
	TRandom3 gaus;

	// 滤波前数据
    vector<float> ori_fx = input_hitsEvent.getFx();
	vector<float> ori_fy = input_hitsEvent.getFy();
	vector<float> ori_fz = input_hitsEvent.getFz();
	vector<float> ori_fe = input_hitsEvent.getFe();
	
	// 滤波后数据
	vector<float> KF_x;
	vector<float> KF_y;
	vector<float> KF_z;

	float e_energy = input_hitsEvent.getTotalEnergy();
	int num_of_oriHits = input_hitsEvent.getNumOfHits();
	cout << "num_of_oriHits = " << num_of_oriHits << endl;
	
	//速度信息产生办法
	vector<float> dLx, dLy, dLz, dL;
	vector<float> speedx_group, speedy_group, speedz_group;
	float dz = 1;
	float thetarnd0 = 13600. / e_energy * sqrt(dz/23520.); //多次散射角
	cout<<"thetarnd0: "<<thetarnd0 <<", e_energy :"<< e_energy<<endl;
	for(int k=0;k<ori_fx.size()-1;k++)
	{	
		float speedx,speedy,speedz=0.;
		
		dLx.push_back(ori_fx[k+1]-ori_fx[k]);
		dLy.push_back(ori_fy[k+1]-ori_fy[k]);
		dLz.push_back(ori_fz[k+1]-ori_fz[k]);
		dL.push_back(sqrt(pow(dLx[k],2)+pow(dLy[k],2)+pow(dLz[k],2)));

		speedx=(gaus.Gaus(dLx[k],thetarnd0));
		speedy=(gaus.Gaus(dLy[k],thetarnd0));
		speedz=(gaus.Gaus(dLz[k],thetarnd0));

		speedx_group.push_back(speedx/sqrt(pow(speedx, 2)+pow(speedy, 2)+pow(speedz, 2)));
		speedy_group.push_back(speedy/sqrt(pow(speedx, 2)+pow(speedy, 2)+pow(speedz, 2)));
		speedz_group.push_back(speedz/sqrt(pow(speedx, 2)+pow(speedy, 2)+pow(speedz, 2)));
		// cout << " point: " << k << ", position: (" << ori_fx[k] << ", " << ori_fy[k] << ", " << ori_fz[k]
		// << "), vector: (" << speedx_group[k] << ", " << speedy_group[k] << ", " << speedz_group[k] << ")" <<endl;
	}
	// cout << "point: " << 0 << ", position: (" << ori_fx[0] << ", " << ori_fy[0] << ", " << ori_fz[0]
	// 	<< "), vector: (" << speedx_group[0] << ", " << speedy_group[0] << ", " << speedz_group[0] << ")" <<endl;

	//参数准备	
	double Speedx, Speedy, Speedz;

	//卡尔曼滤波点和平滑过程点
	double KFx[num_of_oriHits], KFy[num_of_oriHits], KFz[num_of_oriHits], 
		KFvx[num_of_oriHits], KFvy[num_of_oriHits], KFvz[num_of_oriHits], 
		SRx[num_of_oriHits], SRy[num_of_oriHits], SRz[num_of_oriHits], 
		SRvx[num_of_oriHits], SRvy[num_of_oriHits], SRvz[num_of_oriHits],
	    NUM[num_of_oriHits];
		
	SVector3 XMscoli;
	SVector6 Expectation;//6D的预测vector
	SVector6 *ExpectationRecord = new SVector6[num_of_oriHits];
	SVector6 Kalmanresult;//滤波
	SVector6 SmoothResult;//平滑

	SMatrix66 F;//6*6的状态矩阵F
	SMatrix36 measureMatrix;//3*6的测量矩阵H
	measureMatrix(0,0) = 1.0;//H的初始值
	measureMatrix(1,1) = 1.0;
	measureMatrix(2,2) = 1.0;
	SMatrix66 I;//6*6的单位阵
	I(0,0) = 1.0;
	I(1,1) = 1.0;
	I(2,2) = 1.0;
	I(3,3) = 1.0;
	I(4,4) = 1.0;
	I(5,5) = 1.0;
	
	SMatrix66 Priori;
	SMatrix66 *PrioriRecord = new SMatrix66[num_of_oriHits];
	SMatrix66 Posteriori;
	SMatrix66 *PosterioriRecord = new SMatrix66[num_of_oriHits];
	SMatrix63 KalmanGain;
	SMatrix66 SmoothGain;
	SMatrix66 NoiseThMatrix;
	SMatrix33 NoiseMsMatrix;
	SMatrix33 KIn;

	//参数初始化
	Posteriori(0,0)=1.0;
	Posteriori(1,1)=1.0;
	Posteriori(2,2)=1.0;
	Posteriori(3,3)=1.0;
	Posteriori(4,4)=1.0;
	Posteriori(5,5)=1.0;
	
	PosterioriRecord[0] = Posteriori;
	
	KFx[0] = ori_fx[0];
	KFy[0] = ori_fy[0];
	KFz[0] = ori_fz[0];
	KFvx[0] = speedx_group[0];//speedx_group[0]
	KFvy[0] = speedy_group[0];//speedy_group[0]
	KFvz[0] = speedz_group[0];//speedz_group[0]	

	//卡尔曼滤波过程for循环
	for(int n2=1; n2 < num_of_oriHits; n2++)
	{
		Speedx = speedx_group[n2 - 1];
		Speedy = speedy_group[n2 - 1];
		Speedz = speedz_group[n2 - 1];
	
		Kalmanresult[0] = KFx[n2 - 1];
		Kalmanresult[1] = KFy[n2 - 1];
		Kalmanresult[2] = KFz[n2 - 1];
		Kalmanresult[3] = KFvx[n2 - 1];
		Kalmanresult[4] = KFvy[n2 - 1];
		Kalmanresult[5] = KFvz[n2 - 1];
	
		F(0,0) = 1.0;
		F(0,3) = dLz[n2-1] / Speedz;//dt 
		F(1,1) = 1.0;
		F(1,4) = dLz[n2-1] / Speedz;
		F(2,2) = 1.0;
		F(2,5) = dLz[n2-1] / Speedz;
		F(3,3) = 1.0;
		F(4,4) = 1.0;
		F(5,5) = 1.0;
	
		Expectation = F * Kalmanresult;

		ExpectationRecord[n2] = Expectation;//x(k+1)=x(k)+vx*dt

		// float thetarnd0 = 13600. / e_energy * sqrt(DLz[n2-1]/23520.);
		NoiseThMatrix(0,0) = 0.0;
		NoiseThMatrix(1,1) = 0.0;
		NoiseThMatrix(2,2) = 0.0;
		NoiseThMatrix(3,3) = 2.0 * sin(thetarnd0/2) * sqrt(0.5 * (pow(Speedy, 2) + pow(Speedz, 2)));
		NoiseThMatrix(4,4) = 2.0 * sin(thetarnd0/2) * sqrt(0.5 * (pow(Speedx, 2) + pow(Speedz, 2)));
		NoiseThMatrix(5,5) = 2.0 * sin(thetarnd0/2) * sqrt(0.5 * (pow(Speedx, 2) + pow(Speedy, 2)));

		Priori = F * Posteriori*ROOT::Math::Transpose(F) + NoiseThMatrix * NoiseThMatrix;

		PrioriRecord[n2] = Priori;//every state save in PrioriRecord 

		NoiseMsMatrix(0,0) = MeasureErrorx;
		NoiseMsMatrix(1,1) = MeasureErrory;
		NoiseMsMatrix(2,2) = MeasureError;

		//cout<<"NoiseMsMatrix2: "<<NoiseMsMatrix * NoiseMsMatrix<<endl;

		KIn = measureMatrix * Priori * ROOT::Math::Transpose(measureMatrix) + NoiseMsMatrix * NoiseMsMatrix;

		KIn.Invert();

		KalmanGain = Priori * ROOT::Math::Transpose(measureMatrix) * KIn;

		XMscoli[0] = ori_fx[n2];
		XMscoli[1] = ori_fy[n2];
		XMscoli[2] = ori_fz[n2];
	
		Kalmanresult = Expectation + KalmanGain * (XMscoli - measureMatrix * Expectation);//Kalman result
	
		Posteriori = (I - KalmanGain * measureMatrix) * Priori;
	
		PosterioriRecord[n2] = Posteriori;

		KFx[n2] = Kalmanresult[0];
		KFy[n2] = Kalmanresult[1];
		KFz[n2] = Kalmanresult[2];
		KFvx[n2] = Kalmanresult[3];
		KFvy[n2] = Kalmanresult[4];
		KFvz[n2] = Kalmanresult[5];

		NUM[n2] = n2;
	}
	//卡尔曼滤波过程结束

	//平滑过程开始
	SmoothResult[0] = KFx[num_of_oriHits - 1];
	SmoothResult[1] = KFy[num_of_oriHits - 1];
	SmoothResult[2] = KFz[num_of_oriHits - 1];
	SmoothResult[3] = KFvx[num_of_oriHits - 1];
	SmoothResult[4] = KFvy[num_of_oriHits - 1];
	SmoothResult[5] = KFvz[num_of_oriHits - 1];
	SRx[num_of_oriHits - 1] = KFx[num_of_oriHits - 1];
	SRy[num_of_oriHits - 1] = KFy[num_of_oriHits - 1];
	SRz[num_of_oriHits - 1] = KFz[num_of_oriHits - 1];
	for(int n3=num_of_oriHits-2; n3 > -1; n3--)
	{
		Kalmanresult[0] = KFx[n3];
		Kalmanresult[1] = KFy[n3];
		Kalmanresult[2] = KFz[n3];
		Kalmanresult[3] = KFvx[n3];
		Kalmanresult[4] = KFvy[n3];
		Kalmanresult[5] = KFvz[n3];

		PrioriRecord[n3 + 1].Invert();
	
		F(0,0) = 1.0;
		F(0,3) = dLz[n3] / speedz_group[n3];
		F(1,1) = 1.0;
		F(1,4) = dLz[n3] / speedz_group[n3];
		F(2,2) = 1.0;
		F(2,5) = dLz[n3] / speedz_group[n3];
		F(3,3) = 1.0;
		F(4,4) = 1.0;
		F(5,5) = 1.0;
	
		SmoothGain = PosterioriRecord[n3] * ROOT::Math::Transpose(F) * PrioriRecord[n3 + 1];

		SmoothResult = Kalmanresult + SmoothGain * (SmoothResult - ExpectationRecord[n3 + 1]);

		SRx[n3] = SmoothResult[0];
		SRy[n3] = SmoothResult[1];
		SRz[n3] = SmoothResult[2];
		SRvx[n3] = SmoothResult[3];
		SRvy[n3] = SmoothResult[4];
		SRvz[n3] = SmoothResult[5];
	
	}
	//平滑过程结束，读出数据
	
	for (int k=0; k < sizeof(SRx)/sizeof(SRx[0]); k++)
	{	
		float nol_SRv[3];
		nol_SRv[0] = SRvx[k];
		nol_SRv[1] = SRvy[k];
		nol_SRv[2] = SRvz[k];

		float mag_a = 0;
		for (int l=0; l < 3; l++)
		{
			mag_a += pow(nol_SRv[l], 2);
		}
		for (int h=0; h < 3; h++)
		{
			nol_SRv[h] = nol_SRv[h]/mag_a;
		}

		KF_x.push_back(SRx[k]);
		KF_y.push_back(SRy[k]);
		KF_z.push_back(SRz[k]);
		// cout << " SRdata: " << k << ", position: (" << SRx[k] << ", " << SRy[k] << ", " << SRz[k]
		// << "), vector: (" << nol_SRv[0] << ", " << nol_SRv[1] << ", " << nol_SRv[2] << ")" <<endl;
	}

	delete[] ExpectationRecord;
	delete[] PrioriRecord;
	delete[] PosterioriRecord;

	ori_fx.clear();
	ori_fy.clear();
	ori_fz.clear();
	dLx.clear();
	dLy.clear();
	dLz.clear();
	speedx_group.clear();
	speedy_group.clear();
	speedz_group.clear();

	HitsEvent KF_result(KF_x, KF_y, KF_z, ori_fe);
	return KF_result;
}